#!/usr/bin/env python
# -*- coding:utf-8 _*-

from __future__ import print_function

import os
import sys
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
import matplotlib.pyplot as plt

analyPath = os.getcwd()
plt.rcParams["font.family"] = "Times New Roman"
plt.rcParams["font.size"] = 10.5
dg = int(sys.argv[3]) if len(sys.argv) > 3 else 3


class Train_CV_SGBRT(object):

    def __init__(self):
        self.data_path = analyPath + "/data"

    def analy(self, bench, p):

        # 读取数据
        data = pd.read_csv(os.path.join(self.data_path, bench + ".csv"))
        data = pd.DataFrame(data[[p, 'result']].values)
        sorts = list(set(data[0]))
        sorts.sort()

        # 定义线性模型
        Linear = LinearRegression()
        degree = min(len(sorts) - 1, dg)
        quadratic_featurizer = PolynomialFeatures(degree=degree)

        # 拟合并做出预测
        sorts = list(set(data[0]))
        sorts.sort()
        X = quadratic_featurizer.fit_transform(data[[0]])
        Linear.fit(X, data[1])
        x = np.arange(min(sorts), max(sorts), (max(sorts) - min(sorts)) / 50)
        X = quadratic_featurizer.fit_transform(x.reshape(-1, 1))
        y = Linear.predict(X)

        # 画图
        plt.figure(figsize=(4, 3), frameon=False)
        plt.plot(x, y)
        plt.xlabel(p)
        plt.tight_layout()
        plt.savefig(analyPath + '/result/viewimportance/' + bench + '-' + p + '-' + str(dg) + '.pdf')


if __name__ == '__main__':
    train_cv_sgbrt = Train_CV_SGBRT()
    train_cv_sgbrt.__init__()
    train_cv_sgbrt.analy(sys.argv[1], sys.argv[2])
